'''
A quick and dirty VTK based viewer for volume data.
'''


import numpy, wx, vtk, defacing ,thread

try:
    from transfer_function import TransferFunctionWidget
except ImportError:
    print 'You need to get the TransferFunction from: "http://www.siafoo.net/snippet/122"'
    exit(1)

try:
    from wxVTKRenderWindowInteractor import wxVTKRenderWindowInteractor
except ImportError:
    print 'Warning: If you are getting flickering get wxVTKRenderWindowInteractor from: "http://www.siafoo.net/snippet/312"'
    from vtk.wx.wxVTKRenderWindowInteractor import wxVTKRenderWindowInteractor

try:       
    from vtkImageImportFromArray import vtkImageImportFromArray
except ImportError:
    print 'You need to get the updated vtkImageImportFromArray from: "http://www.siafoo.net/snippet/313" '
    exit(1)
    

class TransferGraph(wx.Dialog):

    def __init__(self, parent, id=wx.ID_ANY, title="", pos=wx.DefaultPosition,
                 size=wx.DefaultSize, style=wx.DEFAULT_FRAME_STYLE):

        wx.Dialog.__init__(self, parent, id, title, pos, size, style)
        self.mainPanel = wx.Panel(self, -1)

        # Create some CustomCheckBoxes
        self.t_function = TransferFunctionWidget(self.mainPanel, -1, "", size=wx.Size(300, 150))
       
        # Layout the items with sizers
        mainSizer = wx.BoxSizer(wx.VERTICAL)
        mainSizer.Add(self.mainPanel, 1, wx.EXPAND)
        
        self.SetSizer(mainSizer)
        mainSizer.Layout()
        
    def get_function(self):
        return self.t_function

class VolumeRenderHeadVTK(wxVTKRenderWindowInteractor): 
    def __init__(self, parent, data_set,thres):
	wxVTKRenderWindowInteractor.__init__(self, parent, -1, size=parent.GetSize())

	ren = vtk.vtkRenderer()
	
	self.GetRenderWindow().AddRenderer(ren)
	
	img = self.LoadVolumeData(data_set,thres)
	
	pix_diag = 5.0
	
	#volMapper = vtk.vtkVolumeTextureMapper3D()
	
	volMapper = vtk.vtkVolumeRayCastMapper()
	compositeFunction = vtk.vtkVolumeRayCastCompositeFunction()
	compositeFunction.SetCompositeMethodToInterpolateFirst()
	volMapper.SetVolumeRayCastFunction(compositeFunction)
	volMapper.SetSampleDistance(pix_diag / 5.0)
	#er.CroppingOn()
	volMapper.SetCroppingRegionPlanes(00.0,10000.0,0.0,100.0,0.0,10000.0)
	volMapper.SetCroppingRegionFlagsToSubVolume()
	volMapper.SetInputConnection(img.GetOutputPort())

        # Transfer Functions
        self.opacity_tf = vtk.vtkPiecewiseFunction()
        self.color_tf = vtk.vtkColorTransferFunction()

        # The property describes how the data will look
        self.volProperty = volProperty = vtk.vtkVolumeProperty()
        volProperty.SetColor(self.color_tf)
        volProperty.SetScalarOpacity(self.opacity_tf)
        volProperty.ShadeOn()
        volProperty.SetInterpolationTypeToLinear()
        volProperty.SetScalarOpacityUnitDistance(pix_diag)

        vol = vtk.vtkVolume()
        vol.SetMapper(volMapper)
        vol.SetProperty(volProperty)

	ren.AddVolume(vol)

	contours= defacing.ExtractSurface(img)

	
	normals = defacing.ExtractLargest(contours)
	
	

	normals=defacing.SmoothSurface(normals,200)
	
	contMapper = vtk.vtkPolyDataMapper()
	contMapper.SetInputConnection(normals.GetOutputPort())
	contMapper.ScalarVisibilityOff()
	
	contActor = vtk.vtkActor()
	contActor.SetMapper(contMapper)

	contActor.GetProperty().SetDiffuseColor(1, .49, .25)
	#contActor.GetProperty().SetSpecular(.3)
	contActor.GetProperty().SetSpecularPower(20)

	

        #ren.AddActor(contActor)

	#contMapper.SetCroppingRegionFlagsToSubVolume()
	#cutter=vtk.vtkCutter()
	#cutter.SetCutFunction(plane)
	#cutter.SetInput(normals.GetOutput())
	#cutter.Update()

	"""     boxWidget = vtk.vtkBoxWidget()
	boxWidget.SetInteractor(self)
	boxWidget.SetPlaceFactor(1.0)

        # The implicit function vtkPlanes is used in conjunction with the
        # volume ray cast mapper to limit which portion of the volume is
        # volume rende#red.
        planes = vtk.vtkPlanes()
        def ClipVolumeRender(obj, event):
            obj.GetPlanes(planes)
            print planes
            volMapper.SetClippingPlanes(planes)

	# Place the interactor initially. The output of the reader is used to
	# place the box widget.
	boxWidget.SetInput(img.GetOutput())
	boxWidget.TranslationEnabledOn()
	boxWidget.PlaceWidget()
	boxWidget.InsideOutOn()
	boxWidget.OutlineFaceWiresOff()
	boxWidget.AddObserver("InteractionEvent", ClipVolumeRender)

        outline = vtk.vtkOutlineFilter()
        outline.SetInputConnection(normals.GetOutputPort())
        outlineMapper = vtk.vtkPolyDataMapper()
        outlineMapper.SetInputConnection(outline.GetOutputPort())
        outlineActor = vtk.vtkActor()
        outlineActor.SetMapper(outlineMapper)

        outlineProperty = boxWidget.GetOutlineProperty()
        outlineProperty.SetRepresentationToWireframe()
        outlineProperty.SetAmbient(1.0)
        outlineProperty.SetAmbientColor(1, 1, 1)
        outlineProperty.SetLineWidth(3)
        
        selectedOutlineProperty = boxWidget.GetSelectedOutlineProperty()
        selectedOutlineProperty.SetRepresentationToWireframe()
        selectedOutlineProperty.SetAmbient(1.0)
        selectedOutlineProperty.SetAmbientColor(1, 0, 0)
        selectedOutlineProperty.SetLineWidth(3)

	
	ren.AddActor(outlineActor) #Edited by Costas 6/29/2011-> enable to get bounding box"""   
	style = vtk.vtkInteractorStyleTrackballCamera()
	self.SetInteractorStyle(style)
	

        self.t_graph = TransferGraph(self)

        self.SetFocus()
        
        # So we know when new values are added / changed on the tgraph
        self.t_graph.Connect(-1, -1, wx.wxEVT_COMMAND_SLIDER_UPDATED, 
                             self.OnTGraphUpdate)
       
        self.OnTGraphUpdate(None)
        
        self.lighting = True

        # This is the transfer graph
        #self.t_graph.Show()

    def getParser():

	parser = OptionParser()
	parser.add_option("-f", "--filename",dest= "filename",type = "string", default="")
	parser.add_option("-t", "--threshole", dest = "threshole",type = "int",default= 0)
	
	return parser


   

    def OnTGraphUpdate(self, event):

        self.color_tf.RemoveAllPoints()
        self.opacity_tf.RemoveAllPoints()
        
        for p in self.t_graph.t_function.points:
            rgba = p.get_rgba()
            self.color_tf.AddRGBPoint(p.value, rgba[0]/255.0, 
                                 rgba[1]/255.0, rgba[2]/255.0)

            self.opacity_tf.AddPoint(p.value, rgba[3])

        self.Refresh()

    def LoadVolumeData(self, data_set,thres):

        if data_set.endswith('nii') or data_set.endswith('nii.gz'):
            try:
                from nifti import NiftiImage
            except ImportError:
                print "Apparently you don't have PyNIfTI installed, see http://www.siafoo.net/snippet/310 for instructions"
                exit(1)
            nim = NiftiImage(data_set)
            img_data = nim.data
        elif data_set.endswith('hdr'):
            # Use the header to figure out the shape of the data
            # then load the raw data and reshape the array 
            shape = [int(x) for x in open(data_set).readline().split()]
            img_data = numpy.frombuffer(open(data_set.replace('.hdr', '.dat'), 'rb').read(),
                                        numpy.uint8)\
                        .reshape((shape[2], shape[1], shape[0]))
        

	img_data = numpy.cast['uint8'](img_data)               
	img = vtkImageImportFromArray()
	'''print img_data.max()
	for i in range(img_data.shape[0]):
		for j in range(img_data.shape[1]):
			if(img_data[i,j,img_data.shape[2]-1]>-1):
				for k in range(20):
					img_data[i,j,img_data.shape[2]-k-1] =1
	'''			
	#thread.start_new_thread(defacing.FrontOrthoRayCast, (img_data,thres,data_set))

       # img_data = defacing.FullOrthoCast(img_data,thres,True);

	print "starting"
	img.SetArray(img_data) #Edited Costas 6/29/2011d	

        return img


    def OnKeyDown(self, event):
        # Yes, I did just hack this... maybe one day when things work
        # the way the docs claim... hacks won't be necessary.
        
        key_code = event.GetKeyCode()
        
        if key_code == ord('l'):
            self.lighting = not self.lighting
            print 'Lighting', self.lighting 
            self.volProperty.SetShade(self.lighting)
            self.Refresh()

	if key_code == ord('s'):
	    # screenshot code:
	    w2if = vtk.vtkWindowToImageFilter()
	    w2if.SetInputBufferTypeToZBuffer()
	    w2if.SetInput(self.GetRenderWindow())
	    w2if.Update()
	    z = vtk.vtkFloatArray()
	    self.GetRenderWindow().GetZbufferData(0,0,300,300,z)
	    #print self.GetRenderWindow().GetZBufferData()
	    writer = vtk.vtkTIFFWriter()
	    writer.SetFileName("screenshot.tif")
	    writer.SetInput(w2if.GetOutput())
	    print z.GetRange();
	    self.GetRenderWindow().Render()
	    writer.Write()
	else:
	    wxVTKRenderWindowInteractor.OnKeyDown(self, event)


def show(data):
    app = wx.App()
    frame = wx.Frame(None, -1, 'Volume Rendering with VTK', wx.DefaultPosition, wx.Size(600, 600))
    
    

    canvas = VolumeRenderHeadVTK(frame, data_filename,thres)

    frame.Show()
    app.MainLoop()

def main():
    parser = getParser()
    (options, args) = parser.parse_args()
    
    app = wx.App()
    frame = wx.Frame(None, -1, 'Volume Rendering with VTK', wx.DefaultPosition, wx.Size(600, 600))
    
    try:
        data_filename = options.filename
    except:
#        data_filename = 'avg152T1_RL_nifti.nii.gz' # NIfTI file
        data_filename = '../data/MRI_head.hdr'
    try:
        thres = int(options.threshole)
    except:
#        data_filename = 'avg152T1_RL_nifti.nii.gz' # NIfTI file
	thres = 100
	print 'correct usage : program file threshold , setting thres to '+str(thres)

    canvas = VolumeRenderHeadVTK(frame, data_filename,thres)

    frame.Show()
    app.MainLoop()




if __name__ == '__main__':
    main()